package cn.monkey.io.transport.core.reactive;

import cn.monkey.io.transport.core.Exchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.lang.invoke.MethodHandles;
import java.lang.invoke.VarHandle;
import java.util.Collection;
import java.util.Iterator;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public abstract class AbstractReactiveExchangeManager<C, EXCHANGE extends Exchange> implements ReactiveExchangeManager<C, EXCHANGE> {
    protected final ReactiveExchangeFactory<C, EXCHANGE> factory;
    protected volatile Map<String, EXCHANGE> exchangeMap;

    protected static final VarHandle EXCHANGE_MAP_UPDATER;

    static {
        MethodHandles.Lookup lookup = MethodHandles.lookup();
        try {
            EXCHANGE_MAP_UPDATER = lookup.findVarHandle(AbstractReactiveExchangeManager.class, "exchangeMap", Map.class);
        } catch (NoSuchFieldException | IllegalAccessException e) {
            throw new RuntimeException(e);
        }
    }

    public AbstractReactiveExchangeManager(ReactiveExchangeFactory<C, EXCHANGE> factory,
                                           Map<String, EXCHANGE> exchangeMap) {
        this.factory = factory;
        this.exchangeMap = exchangeMap;
    }

    public AbstractReactiveExchangeManager(ReactiveExchangeFactory<C, EXCHANGE> factory) {
        this(factory, new ConcurrentHashMap<>());
    }

    @Override
    public Mono<EXCHANGE> findOrCreate(String key, C c) {
        final Map<String, EXCHANGE> exchangeMap = this.exchangeMap;
        return Mono.justOrEmpty(exchangeMap.get(key))
                .switchIfEmpty(Mono.defer(() -> this.factory.apply(key, null, c)))
                .map(exchange -> {
                    boolean[] success = {false};
                    EXCHANGE e = exchangeMap.compute(key, (k, v) -> {
                        if (v == null || !v.isActive()) {
                            success[0] = true;
                            return exchange;
                        }
                        return v;
                    });
                    if (success[0]) this.exchangeMap = new ConcurrentHashMap<>(exchangeMap);
                    return e;
                });
    }

    protected boolean tryUpdateExchangeMap(Map<String, EXCHANGE> oldMap, Map<String, EXCHANGE> newMap) {
        return EXCHANGE_MAP_UPDATER.compareAndSet(this, oldMap, newMap);
    }

    @Override
    public Mono<Void> destroy(EXCHANGE exchange) {
        final Map<String, EXCHANGE> exchangeMap = this.exchangeMap;
        Iterator<Map.Entry<String, EXCHANGE>> iterator = exchangeMap.entrySet().iterator();
        while (iterator.hasNext()) {
            Map.Entry<String, EXCHANGE> e = iterator.next();
            Exchange value = e.getValue();
            if (value == exchange) {
                iterator.remove();
                this.tryUpdateExchangeMap(exchangeMap, exchangeMap);
                return value.destroy();
            }
        }
        return Mono.error(new IllegalArgumentException("exchange not exists"));
    }
}
